import os
import time

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.keys import Keys

import wget


class WebDriverChrome(object):

    def __init__(self):
        # self.chrome_options = webdriver.ChromeOptions()
        self.chrome_options = Options()
        self.driver = self.StartWebdriver()

    def options(self, download_path):
        self.chrome_options.add_argument('--headless')  # 无头模式
        self.chrome_options.add_argument("start-maximized")
        self.chrome_options.add_argument('--no-sandbox')
        self.chrome_options.add_argument('--disable-gpu')
        self.chrome_options.add_argument('--disable-dev-shm-usage')
        self.chrome_options.add_experimental_option("excludeSwitches", ["enable-automation"])
        self.chrome_options.add_experimental_option("useAutomationExtension", False)
        prefs = {'profile.default_content_settings.popups': 0,  # 防止保存弹窗
                 'download.default_directory': download_path,  # 设置默认下载路径
                 "profile.default_content_setting_values.automatic_downloads": 1  # 允许多文件下载
                 }
        self.chrome_options.add_experimental_option('prefs', prefs)
        return self.chrome_options

    def StartWebdriver(self, download_path='data/cvpr'):
        self.driver = webdriver.Chrome('D:/Download/Google Download/softwares/chromedriver.exe', options=self.options(download_path))
        with open('../source/stealth.min.js') as f:
            js = f.read()
            self.driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", {
                "source": js
            })
        return self.driver


class CVPR_spider():
    def __init__(self):
        self.url = 'https://openaccess.thecvf.com/CVPR'

    def url_year(self, year):
        '''
        返回 self.url + str(year) 拼接的地址
        :param year:
        :return:
        '''
        return self.url + str(year)

    def init_path(self, word):
        '''
        根据关键词和当前时间创建下载目录
        :param word:
        :return:
        '''
        if os.path.exists('data/cvpr/' + word) is False:
            os.makedirs('data/cvpr/' + word)
        now = time.strftime("%Y-%m-%d_%H_%M_%S", time.localtime())
        download_path = 'data/cvpr/' + word + '/' + str(now)
        if os.path.exists('data/cvpr/' + word + '/' + str(now)) is False:
            os.makedirs('data/cvpr/' + word + '/' + str(now))
        return download_path

    def get_pdf(self, keyword, year):
        '''
        下载 pdf
        :param keyword:
        :param year:
        :return:
        '''
        # 初始化
        download_path = self.init_path(keyword)
        self.driver = WebDriverChrome().StartWebdriver(download_path)
        self.driver.get(self.url_year(year))
        # 搜索框输入关键词并搜索
        input_button = self.driver.find_element_by_tag_name('input')
        input_button.send_keys(keyword)
        time.sleep(1)
        input_button.send_keys(Keys.ENTER)

        # 所有的 pdf
        pdf_name_list = [ptitle.text + '.pdf' for ptitle in self.driver.find_elements_by_class_name('ptitle')]
        pdf_links = self.driver.find_elements_by_link_text('pdf')
        for i in range(len(pdf_links)):
            print(type(pdf_links[i].get_attribute('href')), pdf_links[i].get_attribute('href'))
            wget.download(pdf_links[i].get_attribute('href'), out=download_path+'/'+pdf_name_list[i])
            time.sleep(1)

        # 结束
        self.driver.quit()


if __name__ == '__main__':
    cvpr_spider = CVPR_spider()
    cvpr_spider.get_pdf('small object detection', 2021)